"""
few shot demo example
"""
from pprint import pprint
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate

## 定义字典，喂给大模型的参考示例。那么大模型就会按照示例，进行推理回答。
examples = [
    {"question": "谁活的更久，默罕默德·阿里还是艾伦·图灵？",
     "answer": """
     是否需要后续问题：是的。
     后续问题：默罕默德·阿里去世时多大年纪？
     中间答案：默罕默德·阿里去世时74岁。
     后续问题：艾伦·图灵去世时多大年纪？
     中间答案：艾伦·图灵去世时41岁。
     所以最终答案是：默罕默德·阿里
     """
     },

    {"question": "克雷格斯列表的创始人是什么时候出生的？",
     "answer": """
     是否需要后续问题：是的。
     后续问题：克雷格斯列表的创始人是谁？
     中间答案：克雷格斯列表的创始人是克雷格·纽马克。
     后续问题：克雷格·纽马克是什么时候出生的？
     中间答案：克雷格·纽马克于1952年12月16日出生
     所以最终答案是：1952年12月16日
     """
     },

    {"question": "乔治·华盛顿的外祖父是谁？",
     "answer": """
      是否需要后续问题：是的。
      后续问题：乔治·华盛顿的母亲是谁？
      中间答案：乔治·华盛顿的母亲是玛丽·波尔·华盛顿。
      后续问题：玛丽·波尔·华盛顿的父亲是谁？
      中间答案：玛丽·波尔·华盛顿的父亲是约瑟夫·波尔
      所以最终答案是：约瑟夫·波尔
      """
     },

    {"question": "《大白鲨》和《皇家赌场》的导演都来自同一个国家吗？",
     "answer": """
      是否需要后续问题：是的。
      后续问题：《大白鲨》的导演是谁？
      中间答案：《大白鲨》的导演是史蒂文·斯皮尔伯格。
      后续问题：史蒂文·斯皮尔伯格来自哪个国家？
      中间答案：美国
     后续问题：《皇家赌场》的导演是谁？
      中间答案：《皇家赌场》的导演是马丁·坎贝尔。
      后续问题：马丁·坎贝尔来自哪个国家？
      中间答案：新西兰
      所以最终答案是：不是
      """
     }
]

"""
单个示例---生成提示词模板
把dic中的变量，替换成字符串中的{XX}变量
"""
def one_example_demo():
    example_prompt = PromptTemplate.from_template("问题：{question}\n{answer}")
    # 支持传入一个元素
    print(example_prompt.invoke(examples[0]).to_string())

"""
多个示例
"""

examples2 = [
    {"question": "什么是人工智能？", "answer": "人工智能是..."},
    {"question": "机器学习是什么？", "answer": "机器学习是..."}
]
def many_example_demo():
    example_prompt = PromptTemplate.from_template("问题:{question}\n{answer}")
    prompt = FewShotPromptTemplate(
        examples=examples, # 示例字典
        example_prompt= example_prompt, # 提示词模板
        suffix="问题:{input}",
        input_variables=["input"]
    )
    print(
        prompt.invoke({"input": "乔治·华盛顿的父亲是谁？"}).to_string()
    )
    # user_input = "乔治·华盛顿的父亲是谁？"
    # formatted_prompt = prompt.format(input=user_input)
    # print(formatted_prompt)


if __name__ == "__main__":
    # # 设置日志级别 (DEBUG 查看更多细节)
    # logger.setLevel(logging.WARN)
    # 启动应用
    # one_example_demo()
    many_example_demo()